﻿// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;

namespace Microsoft.SemanticKernel.ChatCompletion;

/// <summary>
/// Class sponsor that holds extension methods for <see cref="IChatCompletionService"/> interface.
/// </summary>
public static class ChatCompletionServiceExtensions
{
    /// <summary>
    /// Get chat multiple chat message content choices for the prompt and settings.
    /// </summary>
    /// <remarks>
    /// This should be used when the settings request for more than one choice.
    /// </remarks>
    /// <param name="chatCompletionService">Target chat completion service.</param>
    /// <param name="prompt">The standardized prompt input.</param>
    /// <param name="executionSettings">The AI execution settings (optional).</param>
    /// <param name="kernel">The <see cref="Kernel"/> containing services, plugins, and other state for use throughout the operation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>List of different chat message content choices generated by the remote model</returns>
    public static Task<IReadOnlyList<ChatMessageContent>> GetChatMessageContentsAsync(
        this IChatCompletionService chatCompletionService,
        string prompt,
        PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null,
        CancellationToken cancellationToken = default)
    {
        // Try to parse the text as a chat history
        if (ChatPromptParser.TryParse(prompt, out var chatHistoryFromPrompt))
        {
            return chatCompletionService.GetChatMessageContentsAsync(chatHistoryFromPrompt, executionSettings, kernel, cancellationToken);
        }

        // Otherwise, use the prompt as the chat user message
        var chatHistory = new ChatHistory();
        chatHistory.AddUserMessage(prompt);

        return chatCompletionService.GetChatMessageContentsAsync(chatHistory, executionSettings, kernel, cancellationToken);
    }

    /// <summary>
    /// Get a single chat message content for the prompt and settings.
    /// </summary>
    /// <param name="chatCompletionService">The target <see cref="IChatCompletionService"/> interface to extend.</param>
    /// <param name="prompt">The standardized prompt input.</param>
    /// <param name="executionSettings">The AI execution settings (optional).</param>
    /// <param name="kernel">The <see cref="Kernel"/> containing services, plugins, and other state for use throughout the operation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>Single chat message content generated by the remote model.</returns>
    public static async Task<ChatMessageContent> GetChatMessageContentAsync(
        this IChatCompletionService chatCompletionService,
        string prompt,
        PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null,
        CancellationToken cancellationToken = default)
        => (await chatCompletionService.GetChatMessageContentsAsync(prompt, executionSettings, kernel, cancellationToken).ConfigureAwait(false))
            .Single();

    /// <summary>
    /// Get a single chat message content for the chat history and settings provided.
    /// </summary>
    /// <param name="chatCompletionService">The target <see cref="IChatCompletionService"/> interface to extend.</param>
    /// <param name="chatHistory">The chat history to complete.</param>
    /// <param name="executionSettings">The AI execution settings (optional).</param>
    /// <param name="kernel">The <see cref="Kernel"/> containing services, plugins, and other state for use throughout the operation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>List of different chat results generated by the remote model</returns>
    public static async Task<ChatMessageContent> GetChatMessageContentAsync(
        this IChatCompletionService chatCompletionService,
        ChatHistory chatHistory,
        PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null,
        CancellationToken cancellationToken = default)
        => (await chatCompletionService.GetChatMessageContentsAsync(chatHistory, executionSettings, kernel, cancellationToken).ConfigureAwait(false))
            .Single();

    /// <summary>
    /// Get streaming chat message contents for the chat history provided using the specified settings.
    /// </summary>
    /// <exception cref="NotSupportedException">Throws if the specified type is not the same or fail to cast</exception>
    /// <param name="chatCompletionService">The target <see cref="IChatCompletionService"/> interface to extend.</param>
    /// <param name="prompt">The standardized prompt input.</param>
    /// <param name="executionSettings">The AI execution settings (optional).</param>
    /// <param name="kernel">The <see cref="Kernel"/> containing services, plugins, and other state for use throughout the operation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>Streaming list of different completion streaming string updates generated by the remote model</returns>
    public static IAsyncEnumerable<StreamingChatMessageContent> GetStreamingChatMessageContentsAsync(
        this IChatCompletionService chatCompletionService,
        string prompt,
        PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null,
        CancellationToken cancellationToken = default)
    {
        // Try to parse the text as a chat history
        if (ChatPromptParser.TryParse(prompt, out var chatHistoryFromPrompt))
        {
            return chatCompletionService.GetStreamingChatMessageContentsAsync(chatHistoryFromPrompt, executionSettings, kernel, cancellationToken);
        }

        // Otherwise, use the prompt as the chat user message
        var chatHistory = new ChatHistory();
        chatHistory.AddUserMessage(prompt);

        return chatCompletionService.GetStreamingChatMessageContentsAsync(chatHistory, executionSettings, kernel, cancellationToken);
    }

    /// <summary>Creates an <see cref="IChatClient"/> for the specified <see cref="IChatCompletionService"/>.</summary>
    /// <param name="service">The chat completion service to be represented as a chat client.</param>
    /// <returns>
    /// The <see cref="IChatClient"/>. If the <paramref name="service"/> is an <see cref="IChatClient"/>, the <paramref name="service"/>
    /// will be returned. Otherwise, a new <see cref="IChatClient"/> will be created that wraps <paramref name="service"/>.
    /// </returns>
    [Experimental("SKEXP0001")]
    public static IChatClient AsChatClient(this IChatCompletionService service)
    {
        Verify.NotNull(service);

        return service is IChatClient chatClient ?
            chatClient :
            new ChatCompletionServiceChatClient(service);
    }

    /// <summary>Creates an <see cref="IChatCompletionService"/> for the specified <see cref="IChatClient"/>.</summary>
    /// <param name="client">The chat client to be represented as a chat completion service.</param>
    /// <param name="serviceProvider">An optional <see cref="IServiceProvider"/> that can be used to resolve services to use in the instance.</param>
    /// <returns>
    /// The <see cref="IChatCompletionService"/>. If <paramref name="client"/> is an <see cref="IChatCompletionService"/>, <paramref name="client"/> will
    /// be returned. Otherwise, a new <see cref="IChatCompletionService"/> will be created that wraps <paramref name="client"/>.
    /// </returns>
    [Experimental("SKEXP0001")]
    public static IChatCompletionService AsChatCompletionService(this IChatClient client, IServiceProvider? serviceProvider = null)
    {
        Verify.NotNull(client);

        return client is IChatCompletionService chatCompletionService ?
            chatCompletionService :
            new ChatClientChatCompletionService(client, serviceProvider);
    }

    /// <summary>Converts a <see cref="ChatMessageContent"/> to a <see cref="ChatMessage"/>.</summary>
    /// <remarks>This conversion should not be necessary once SK eventually adopts the shared content types.</remarks>
    internal static ChatMessage ToChatMessage(ChatMessageContent content)
    {
        ChatMessage message = new()
        {
            AdditionalProperties = content.Metadata is not null ? new(content.Metadata) : null,
            AuthorName = content.AuthorName,
            RawRepresentation = content.InnerContent,
            Role = content.Role.Label is string label ? new ChatRole(label) : ChatRole.User,
        };

        foreach (var item in content.Items)
        {
            AIContent? aiContent = null;
            switch (item)
            {
                case Microsoft.SemanticKernel.TextContent tc:
                    aiContent = new Microsoft.Extensions.AI.TextContent(tc.Text);
                    break;

                case Microsoft.SemanticKernel.ImageContent ic:
                    aiContent =
                        ic.DataUri is not null ? new Microsoft.Extensions.AI.ImageContent(ic.DataUri, ic.MimeType) :
                        ic.Uri is not null ? new Microsoft.Extensions.AI.ImageContent(ic.Uri, ic.MimeType) :
                        null;
                    break;

                case Microsoft.SemanticKernel.AudioContent ac:
                    aiContent =
                        ac.DataUri is not null ? new Microsoft.Extensions.AI.AudioContent(ac.DataUri, ac.MimeType) :
                        ac.Uri is not null ? new Microsoft.Extensions.AI.AudioContent(ac.Uri, ac.MimeType) :
                        null;
                    break;

                case Microsoft.SemanticKernel.BinaryContent bc:
                    aiContent =
                        bc.DataUri is not null ? new Microsoft.Extensions.AI.DataContent(bc.DataUri, bc.MimeType) :
                        bc.Uri is not null ? new Microsoft.Extensions.AI.DataContent(bc.Uri, bc.MimeType) :
                        null;
                    break;

                case Microsoft.SemanticKernel.FunctionCallContent fcc:
                    aiContent = new Microsoft.Extensions.AI.FunctionCallContent(fcc.Id ?? string.Empty, fcc.FunctionName, fcc.Arguments);
                    break;

                case Microsoft.SemanticKernel.FunctionResultContent frc:
                    aiContent = new Microsoft.Extensions.AI.FunctionResultContent(frc.CallId ?? string.Empty, frc.FunctionName ?? string.Empty, frc.Result);
                    break;
            }

            if (aiContent is not null)
            {
                aiContent.RawRepresentation = item.InnerContent;
                aiContent.AdditionalProperties = item.Metadata is not null ? new(item.Metadata) : null;

                message.Contents.Add(aiContent);
            }
        }

        return message;
    }

    /// <summary>Converts a <see cref="ChatMessage"/> to a <see cref="ChatMessageContent"/>.</summary>
    /// <remarks>This conversion should not be necessary once SK eventually adopts the shared content types.</remarks>
    internal static ChatMessageContent ToChatMessageContent(ChatMessage message, Microsoft.Extensions.AI.ChatCompletion? completion = null)
    {
        ChatMessageContent result = new()
        {
            ModelId = completion?.ModelId,
            AuthorName = message.AuthorName,
            InnerContent = completion?.RawRepresentation ?? message.RawRepresentation,
            Metadata = message.AdditionalProperties,
            Role = new AuthorRole(message.Role.Value),
        };

        foreach (AIContent content in message.Contents)
        {
            KernelContent? resultContent = null;
            switch (content)
            {
                case Microsoft.Extensions.AI.TextContent tc:
                    resultContent = new Microsoft.SemanticKernel.TextContent(tc.Text);
                    break;

                case Microsoft.Extensions.AI.ImageContent ic:
                    resultContent = ic.ContainsData ?
                        new Microsoft.SemanticKernel.ImageContent(ic.Uri) :
                        new Microsoft.SemanticKernel.ImageContent(new Uri(ic.Uri));
                    break;

                case Microsoft.Extensions.AI.AudioContent ac:
                    resultContent = ac.ContainsData ?
                        new Microsoft.SemanticKernel.AudioContent(ac.Uri) :
                        new Microsoft.SemanticKernel.AudioContent(new Uri(ac.Uri));
                    break;

                case Microsoft.Extensions.AI.DataContent dc:
                    resultContent = dc.ContainsData ?
                        new Microsoft.SemanticKernel.BinaryContent(dc.Uri) :
                        new Microsoft.SemanticKernel.BinaryContent(new Uri(dc.Uri));
                    break;

                case Microsoft.Extensions.AI.FunctionCallContent fcc:
                    resultContent = new Microsoft.SemanticKernel.FunctionCallContent(fcc.Name, null, fcc.CallId, fcc.Arguments is not null ? new(fcc.Arguments) : null);
                    break;

                case Microsoft.Extensions.AI.FunctionResultContent frc:
                    resultContent = new Microsoft.SemanticKernel.FunctionResultContent(frc.Name, null, frc.CallId, frc.Result);
                    break;
            }

            if (resultContent is not null)
            {
                resultContent.Metadata = content.AdditionalProperties;
                resultContent.InnerContent = content.RawRepresentation;
                resultContent.ModelId = completion?.ModelId;
                result.Items.Add(resultContent);
            }
        }

        return result;
    }

    internal static List<ChatMessage> ToChatMessageList(ChatHistory chatHistory)
        => chatHistory.Select(ToChatMessage).ToList();
}
